'''
https://leetcode.cn/problems/number-of-atoms/
'''
from collections import Counter


class Solution:
    def fill(self, ans: Counter, pre: Counter, atom, time):
        time = time if time else 1
        if atom:
            ans[atom] += time
        if pre:
            ans += Counter({key: item * time for key, item in pre.items()})

    def countOfAtoms(self, formula: str) -> str:
        n = len(formula)
        where = 0

        def process(i) -> Counter:
            atom, time = '', 0
            ans, pre = Counter(), None
            nonlocal where
            # 什么时候要统计？
            #       1.遇到新的原子
            #       2.遇到了(
            # 遇到数字小字母，还不能确定完整 time or atom
            while i < n and formula[i] != ')':
                if formula[i].isdigit():
                    time = time * 10 + int(formula[i])
                elif formula[i].islower():
                    atom += formula[i]
                else:
                    self.fill(ans, pre, atom, time)
                    atom, pre, time = '', None, 0
                    if formula[i] == '(':
                        pre = process(i + 1)
                        i = where
                    else:
                        atom = formula[i]
                i += 1
            self.fill(ans, pre, atom, time)
            where = i
            return ans

        res = process(0)
        print(res)
        s = [k + ('' if res[k] == 1 else str(res[k])) for k in sorted(res.keys())]
        return ''.join(s)


formula = "H20O"
print(Solution().countOfAtoms(formula))
